# Starting Reference: http://nlp.seas.harvard.edu/2018/04/03/attention.html#greedy-decoding
import torch.cuda
import torch.nn as nn
import torch.optim as optim
from datasets import *
from transformer import Transformer

if __name__ == "__main__":
    # 生成数据集
    enc_inputs, dec_inputs, dec_outputs = make_data()
    MyDataSet = MyDataSet(enc_inputs, dec_inputs, dec_outputs)
    # 将数据集封装成dataloader
    loader = Data.DataLoader(MyDataSet, 2, True)
    # 创建Transformer模型实例并放到GPU上
    model = Transformer().cuda()
    # 定义损失函数为交叉熵，忽略填充位的计算
    criterion = nn.CrossEntropyLoss(ignore_index=0)
    # 定义优化器为随机梯度下降
    optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.99)

    # 开始训练循环
    for epoch in range(50):
        # 遍历数据集，每个batch为一个训练单位
        for enc_inputs, dec_inputs, dec_outputs in loader:
            # enc_inputs : [batch_size, src_len]
            # dec_inputs : [batch_size, tgt_len]
            # dec_outputs: [batch_size, tgt_len]
            # 将数据移动到GPU上
            enc_inputs, dec_inputs, dec_outputs = enc_inputs.cuda(), dec_inputs.cuda(), dec_outputs.cuda()
            # 使用模型进行前向传播，得到预测结果和注意力权重矩阵
            outputs, enc_self_attns, dec_self_attns, dec_enc_attns = model(enc_inputs, dec_inputs)
            # 计算损失
            # outputs  [batch_size * tgt_len, tgt_vocab_size]
            # impo: outputs, dec_outputs.view(-1): 将dec_outputs张量展平成一个一维张量，计算交叉熵损失
            loss = criterion(outputs, dec_outputs.view(-1))
            # 打印损失信息
            print('Epoch:', '%04d' % (epoch + 1), 'loss =', '{:.6f}'.format(loss))
            # 梯度清零
            optimizer.zero_grad()
            # 反向传播计算梯度
            loss.backward()
            # 更新模型参数
            optimizer.step()
    # 保存模型
    torch.save(model, 'model.pt')
    print("保存模型")


# NOTE:
#       torch.nn.CrossEntropyLoss 是一个用于多分类问题的损失函数，其主要作用是计算模型输出的概率分布与实际标签的差异，并返回一个标量值作为损失。
#       该损失函数的输入参数包括两个张量：模型的预测值和真实标签。其中，预测值通常是一个张量，其形状为 (batch_size, num_classes)，表示模型对于每个样本属于不同类别的概率分布，而真实标签是一个大小为 (batch_size,) 的张量，其中每个元素的值代表其对应样本的类别编号。
#       在计算过程中，CrossEntropyLoss 首先对预测值应用 softmax 函数，将其转化为概率分布，然后计算预测概率分布和真实标签的交叉熵损失。该函数的输出是一个标量值，表示一个 batch 的平均损失。
#       在使用 CrossEntropyLoss 时，需要注意的是，模型的最后一层不需要应用 softmax 函数，因为该函数已经被集成在损失函数中。此外，模型的输出和真实标签需要是对应的，即同一 batch 中的输出和标签需要按照相同的顺序排列。